gr = require "graphs"
pq = require "priority_queues"

--[[

relax updates a table *tab* where each entry
is an upper bound on the cost of a route from 
one fixed source s to each node.  

Suppose tab contains a s ~~> u estimate d[u], 
and the graph contains an edge u -> v with cost c.  If 
d[u] + cost(u,v) is better than d[v], or if d[v] 
is not yet defined, then we plug it in.  

We also call the supplied function *fix* to fix 
any related structure such as a priority queue 
that may need to fix the position of v, now that 
we have a lower-cost estimate of d[v].

--]]

function relax(u,v,c,tab,fix) 
   local cu = tab[u].cost
   local cv_entry = tab[v]
   if cu and (not(cv_entry) or cu + c < cv_entry.cost)
   then 
      tab[v] = { cost = cu + c,
		 pi = u }
      fix(v)
   end 
end 

--[[

To initialize a *tab*, we just plug in 
0 as the cost s ~~> s.  

--]]

function initialize(g,s)
   local tab = {} 

   tab[s] = { cost = 0 }
   return tab 
end 

-- A location table is 
-- intended to map an item to its 
-- priority queue location 
-- (if the latter non-nil).  

-- Given an array, return a function 
-- that will install an updated location 
-- for each item.  

function loc_tab_make_loc_fn(loc_tab) 
   local function loc(v,new_l) 
      if v then loc_tab[v] = new_l end 
   end
   return loc 
end 

--[[

Given a *tab*, return a preference function.
This preference fn, when given two args a and 
b, should return true if a is preferred,
meaning that the cost of a as estimated 
by *tab* is lower, or if its cost is defined 
and the cost of b is not.  

--]]

function prefer_by_shortest_route(tab) 
   local function pref(a,b) 
      local da = tab[a] 
      local db = tab[b] 
      if da and db 
      then return da.cost<db.cost
      else return da 
      end 
   end 
   return pref
end 

--[[

Return a triple containing

1.  a heap, representing a priority queue of nodes with a finite
    current bound on their distance,

2.  a distance table recording the currently best known bounds, and 

3.  a location table, mapping each node to its current location 
    in the heap (1), used to find nodes when the need to be heapified-up 
    in the course of reducing their distance bound.  Value is nil if node
    is not currently in the heap.  

Initially, the heap (1) contains only the starting node s.  

--]]

function make_dist_tab(g,s) 
   local tab = initialize(g,s)
   local loc_tab = {} 
   local locate = loc_tab_make_loc_fn(loc_tab)
   local pref = prefer_by_shortest_route(tab)

   local h = pq.build_heap({}, pref, locate)
   pq.insert(h,s)
   return { hp = h, dtab = tab, lt = loc_tab }
end 

--[[

If the heap dt.hp is empty, return false.  
Otherwise, extract the minimal node and return that.  
When removing an element, relax all nodes to which it 
has an edge.  

Rationale:  the distance estimate of the minimal node 
is now final.  Any neighbors of it should get an estimate 
reflecting this final shortest  route to it.  

When relaxing a node v, we need to fix its position 
in the heap, since its distance estimate may have improved 
(decreased).  To fix it, we must either:

   heapify it up if it is already in the heap, 
or insert it into the heap if it is not yet there.  

--]]

function get_minimal(g,dt)

	local min = pq.extract_top(dt.hp) 
   
	if not(min) then 
		return false 
	end
	
	relax( g[min].source, g[min].dest, g[min].cost, pq.dtab, function(v)
		insert( dt.hp, min )
	end )

-- Fill in the blank:
-- Give a definition of this function 
-- using the textbook, the comments in 
-- this file, and the fairly similar 
-- example in topological_sort.lua

end 

-- make a distance table with s at dist 0.
-- Then get_minimal till the 
-- heap is empty.  

function dijkstra_shortest_paths(g,s) 
   local dt = make_dist_tab(g,s)

-- Fill in the blank:
-- Give the rest of the definition of this function.
-- Can be very short.  

   return dt.dtab 
end 

--[[

To print out a route determined from the 
distance table dtab, print out the route to 
the predecessor of the destination, and then 
print the destination.  Use the printer function 
if supplied, otherwise the identity. 

--]] 

function print_route(dtab,dest,printer)
   local p = printer or function(n) return n end 
   local function print_em(last) 
      if last 
      then 
	 print_em(dtab[last].pi)
	 io.write(string.format("  %s\n", p(last)))
      end 
   end 

   io.write(string.format("route to %s, at a total cost of %d\ntraveling via\n",
			  p(dest), dtab[dest].cost))
   print_em(dest)
end 

-- Fill in the blank:  
-- Use examples from cities.lua to construct 
-- shortest routes (relative to mileage) 
-- for several start/end city pairs.  

